import pandas as pd
import pymysql
import json
import datetime
from dateutil.relativedelta import relativedelta


# 每个城市群的主要污染物
def get_db(sql, args=()):
    connect = pymysql.connect(
        host="10.10.4.28",
        database="chinavis_data",
        user="chinavis",
        password="123456",
        port=3306,
        charset='utf8'
    )
    # 获取游标
    cursor = connect.cursor()
    cursor.execute(sql, args)
    connect.commit()
    rv = cursor.fetchall()
    columns = [col[0] for col in cursor.description]
    # 关闭连接
    cursor.close()
    connect.close()
    # return rv
    return [
        dict(zip(columns, row))
        for row in rv
    ]


def get_max(a1, a2, a3, a4, a5, a6):
    a_max = max(a1, a2, a3, a4, a5, a6)
    if a1 == a_max:
        return "PM2.5"
    elif a2 == a_max:
        return "PM10"
    elif a3 == a_max:
        return "SO2"
    elif a4 == a_max:
        return "NO2"
    elif a5 == a_max:
        return "CO"
    elif a6 == a_max:
        return "O3"


def cal_IAQI(begin, type, main_index):
    category = ['长江中游城市群', '哈长城市群', '成渝城市群', '长江三角洲城市群', '中原城市群',
                '北部湾城市群', '关中平原城市群', '呼包鄂榆城市群', '兰西城市群', '粤港澳大湾区', '京津冀城市群',
                '辽中南城市群', '山东半岛城市群', '海峡西岸城市群']
    date_dict = {}
    date_dict['date'] = begin
    sql = "select * from pollution_iaqi_" + type + " where date='" + begin + "'"
    a = get_db(sql)
    columns = a[0].keys()
    data = pd.DataFrame(columns=columns, data=a)

    sql_group = "select * from cluster_" + type + " where date='" + begin + "' and main_index='" + main_index + "'"
    b = get_db(sql_group)
    columns1 = b[0].keys()
    data1 = pd.DataFrame(columns=columns1, data=b)
    all_data = pd.merge(data, data1, how='left', left_on='city_name', right_on='city_name')
    all_data = all_data.groupby('category')
    for c in all_data:
        a_mean = c[1].mean()
        main_p = get_max(a_mean['PM2.5(微克每立方米)'], a_mean['PM10(微克每立方米)'], a_mean['SO2(微克每立方米)'], a_mean['NO2(微克每立方米)'], a_mean['CO(毫克每立方米)'], a_mean['O3(微克每立方米)'])
        date_dict[category[int(c[0])]] = main_p
    return date_dict




def get_weekdata(begin_date, end_date, type, main_index):
    main_pollution = []
    begin = datetime.datetime.strptime(begin_date, '%Y-%m-%d').date()
    end = datetime.datetime.strptime(end_date, '%Y-%m-%d').date()
    week = datetime.timedelta(days=7)
    while 1:
        da_dict = cal_IAQI(begin.strftime('%Y-%m-%d'), type, main_index)
        main_pollution.append(da_dict)
        begin = begin + week
        if begin == end + week:
            break
    return main_pollution

def get_monthdata(begin_date, end_date, type, main_index):
    main_pollution = []
    begin = datetime.datetime.strptime(begin_date, '%Y-%m').date()
    end = datetime.datetime.strptime(end_date, '%Y-%m').date()
    month = relativedelta(months=1)
    while 1:
        test = begin.strftime('%Y-%m-%d')
        list1 = test.split("-")
        begin_d = list1[0] + "-" + list1[1]
        da_dict = cal_IAQI(begin_d, type, main_index)
        main_pollution.append(da_dict)
        if begin == end:
            break
        begin = begin + month
    return main_pollution

def get_quarterdata(begin_date, end_date, type, main_index):
    main_pollution = []
    begin = datetime.datetime.strptime(begin_date, '%Y-%m').date()
    end = datetime.datetime.strptime(end_date, '%Y-%m').date()
    month = relativedelta(months=3)
    while 1:
        test = begin.strftime('%Y-%m-%d')
        list1 = test.split("-")
        begin_d = list1[0] + "-" + list1[1]
        da_dict = cal_IAQI(begin_d, type, main_index)
        main_pollution.append(da_dict)
        if begin == end:
            break
        begin = begin + month
    return main_pollution


def get_main(begin_date, end_date, type, main_index):
    if type == "week":
        main_p = get_weekdata(begin_date, end_date, type, main_index)
        # with open('week_main1.json', 'w', encoding='utf-8') as file:
        #     json.dump(main_p, file, ensure_ascii=False)
    elif type == "month":
        main_p = get_monthdata(begin_date, end_date, type, main_index)
        # with open('month_main1.json', 'w', encoding='utf-8') as file:
        #     json.dump(main_p, file, ensure_ascii=False)
    else:
        main_p = get_quarterdata(begin_date, end_date, type, main_index)
        # with open('quarter_main1.json', 'w', encoding='utf-8') as file:
        #     json.dump(main_p, file, ensure_ascii=False)
    return main_p